import torch

try:
    # Load the saved model
    model = torch.jit.load('C:\\Users\\jxk33\\Desktop\\code\\coding\\homework\\AI_lab\\model_jit.pt')

    # Set input shape and create sample data
    input_shape = [8, 3, 32, 32]
    input_data = torch.randn(input_shape)

    # Ensure model is on CPU and in eval mode
    model = model.cpu()
    model.eval()

    # Trace the model
    traced_model = torch.jit.trace(model, input_data)

    # Save the traced model
    traced_model_name = "traced_model.pt"
    traced_model.save(traced_model_name)

    print(f"Successfully saved traced model to {traced_model_name}")

except Exception as e:
    print(f"Error occurred: {str(e)}")